
#include "acl/acl.h"
#include <iostream>
#include "data_utils.h"
#include "pool/host/pool_host.h"

void EdgeCustomFloat_512_750(uint32_t blockDim, void* stream,  uint8_t* a,  uint8_t* c, ATVC::PoolParam  param);
void EdgeCustomFloat_96_640(uint32_t blockDim, void* stream, uint8_t* a, uint8_t* c, ATVC::PoolParam param);

using OP_TRAITS = ATVC::OpTraits<ATVC::OpInputs<float>, ATVC::OpOutputs<float>, ATVC::OpTemps<>>;;
int32_t main(int32_t argc, char* argv[])
{ // 1: w 2: h 5: enProfiling
    uint32_t totalW = std::stoi(std::string(argv[2]));
    uint32_t totalH = std::stoi(std::string(argv[1]));
    uint32_t tileW = std::stoi(std::string(argv[3]));
    uint32_t tileH = std::stoi(std::string(argv[4]));

    int32_t eleNum = totalW * totalH;
    ATVC::Layout2Dim totalLayout = {totalW, totalH};
    ATVC::Layout2Dim tileLayout = {tileW, tileH};
    printf("Pool start: totalW: %u totalH: %u tileW %u tileH: %u.", totalW, totalH, tileW, tileH);
    bool enableProf = std::string(argv[5]) == "1";
	uint8_t* aHost;
	uint8_t* zHost;
	size_t aByteSize = eleNum * sizeof(float);

	size_t zByteSize = eleNum * sizeof(float);


    ATVC::PoolParam param;
    if (!ATVC::Host::CalcPoolTiling<OP_TRAITS>(totalLayout, tileLayout, param)) {
        printf("Pool tiling error.");
        return -1;
    }
    CHECK_ACL(aclInit({}));
    aclrtContext context;
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    CHECK_ACL(aclrtCreateContext(&context, deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));
	uint8_t* aDevice;
	uint8_t* zDevice;
	CHECK_ACL(aclrtMallocHost((void**)(&aHost), aByteSize));
	CHECK_ACL(aclrtMalloc((void**)&aDevice, aByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
	ReadFile("./input/input_a.bin", aByteSize, aHost, aByteSize);
	CHECK_ACL(aclrtMemcpy(aDevice, aByteSize, aHost, aByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
	CHECK_ACL(aclrtMallocHost((void**)(&zHost), zByteSize));
	CHECK_ACL(aclrtMalloc((void**)&zDevice, zByteSize, ACL_MEM_MALLOC_HUGE_FIRST));


int loopCnt = 1; if(enableProf) {loopCnt = 20;};
for (int i = 0; i < loopCnt; ++i) {
    if (totalW == 512 && totalH == 750) {
        EdgeCustomFloat_512_750(param.tilingData.blockNum, stream, aDevice, zDevice, param);
    } else if (totalW == 96 && totalH == 640) {
        EdgeCustomFloat_96_640(param.tilingData.blockNum, stream, aDevice, zDevice, param);
    }
	
}
	CHECK_ACL(aclrtSynchronizeStream(stream));

	CHECK_ACL(aclrtMemcpy(zHost, zByteSize, zDevice, zByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
	WriteFile("./output/output_z.bin", zHost, zByteSize);

	CHECK_ACL(aclrtFree(aDevice));
	CHECK_ACL(aclrtFreeHost(aHost));

	CHECK_ACL(aclrtFree(zDevice));
	CHECK_ACL(aclrtFreeHost(zHost));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtDestroyContext(context));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());

    return 0;
}
